import heapq
from graphviz import Graph
import networkx as nx

from heapq import heappush, heappop
import networkx as nx
import matplotlib.pyplot as plt

import shutil
import matplotlib

shutil.rmtree(matplotlib.get_cachedir())
def dijkstra_animation(graph, start, end):
    # 创建一个有向图
    G = nx.DiGraph()

    # 添加节点
    for node in graph.keys():
        G.add_node(node)

    # 添加边和边权重
    for node in graph.keys():
        for neighbor, weight in graph[node].items():
            G.add_edge(node, neighbor, weight=weight)

    # 设置起点和终点
    start_node = start
    end_node = end

    # 初始化距离字典和前驱节点字典
    distance = {node: float('inf') for node in graph}
    distance[start_node] = 0
    predecessor = {node: None for node in graph}

    # 初始化堆队列
    heap = [(0, start_node)]

    # 创建图形对象
    fig = plt.figure(figsize=(8, 8))
    plt.title("Dijkstra's Algorithm Animation")
    plt.axis('off')

    # 开始执行Dijkstra算法
    while heap:
        # 从堆队列中获取当前节点
        curr_dist, curr_node = heappop(heap)
        # 如果当前节点已经是终点，则跳出循环
        if curr_node == end_node:
            break
        # 遍历当前节点的邻居节点
        for neighbor, weight in graph[curr_node].items():
            # 计算邻居节点的新距离
            new_dist = curr_dist + weight
            # 如果新距离比原距离小，则更新距离字典和前驱节点字典
            if new_dist < distance[neighbor]:
                distance[neighbor] = new_dist
                predecessor[neighbor] = curr_node
                # 将邻居节点加入堆队列
                heappush(heap, (new_dist, neighbor))

        # 绘制当前图形
        plt.clf()
        pos = nx.spring_layout(G)
        nx.draw_networkx_nodes(G, pos, node_color='lightblue')
        nx.draw_networkx_edges(G, pos, edge_color='gray')
        nx.draw_networkx_labels(G, pos)
        nx.draw_networkx_edge_labels(G, pos, edge_labels=nx.get_edge_attributes(G, 'weight'))
        nx.draw_networkx_nodes(G, pos, nodelist=[curr_node], node_color='red')
        plt.pause(1)

    # 获取最短路径
    path = []
    curr_node = end_node
    while curr_node != start_node:
        path.append(curr_node)
        curr_node = predecessor[curr_node]
    path.append(start_node)
    path.reverse()

    # 绘制最短路径
    plt.clf()
    pos = nx.spring_layout(G)
    nx.draw_networkx_nodes(G, pos, node_color='lightblue')
    nx.draw_networkx_edges(G, pos, edge_color='gray')
    nx.draw_networkx_labels(G, pos)
    nx.draw_networkx_edge_labels(G, pos, edge_labels=nx.get_edge_attributes(G, 'weight'))
    nx.draw_networkx
# 测试代码
graph = {
    '北京': {'石家庄': 310, '郑州': 695, '天津': 120},
    '石家庄': {'北京': 310, '郑州': 450},
    '郑州': {'北京': 695, '石家庄': 450, '广州': 1130, '香港': 1420},
    '天津': {'北京': 120},
    '广州': {'郑州': 1130, '香港': 115},
    '香港': {'郑州': 1420, '广州': 115}
}

if __name__ == '__main__':
    dijkstra_animation(graph, '北京', '香港')
